# -*- coding: utf-8 -*-
# @File:     sentence_sorted.py
# @Author:
# @DateTime: 2025/10/23/16:09

from pathlib import Path
current_path = Path(__file__).resolve().parent
from sentence_transformers import CrossEncoder

RERANK_MODEL = Path.joinpath(current_path,  r'sentence_models\models--cross-encoder--ms-marco-MiniLM-L6-v2\snapshots\c5ee24cb16019beea0893ab7796b1df96625c6b8')     # cross-encoder/ms-marco-MiniLM-L6-v2


def docs_scores_sorted(docs, query_content):
    """
    调用re-rank模型对获取的结果再进行一次相似度计算及排序
    :param query_content:
    :param docs:
    :param query_text:
    :return:
    """
    model = CrossEncoder(RERANK_MODEL, max_length=512)  # type: ignore
    scores = model.predict([(query_content, doc) for doc in docs])
    sorted_list = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
    # print(sorted_list)
    documents = list(map(lambda x: x[1], sorted_list))
    # print(documents)
    return documents


if __name__ == '__main__':


    pass